#!/usr/bin/env python3
"""
Neo4j MCP Server using FastMCP

A Model Context Protocol server that provides tools for interacting with Neo4j databases.
"""

import asyncio
import json
import logging
import os
import sys
from typing import Any, Dict, List, Optional

try:
    from mcp.server.fastmcp import FastMCP  # type: ignore
except ImportError:
    print("Error: mcp.server.fastmcp not found. Please install the MCP SDK.")
    sys.exit(1)

try:
    from neo4j import GraphDatabase  # type: ignore
    from neo4j.exceptions import ServiceUnavailable, AuthError  # type: ignore
except ImportError:
    print("Error: neo4j not found. Please install neo4j driver.")
    sys.exit(1)

# Configure logging
log_level = os.getenv("LOG_LEVEL", "INFO").upper()
logging.basicConfig(level=getattr(logging, log_level, logging.INFO))
logger = logging.getLogger(__name__)

# Create FastMCP server
mcp = FastMCP("Neo4j MCP Server")

# Global Neo4j driver
driver = None

def get_neo4j_driver():
    """Get or create Neo4j driver"""
    global driver
    if driver is None:
        try:
            # Get connection parameters from environment variables
            uri = os.getenv("NEO4J_URI", "bolt://172.16.5.194:32263")
            username = os.getenv("NEO4J_USERNAME", "neo4j")
            password = os.getenv("NEO4J_PASSWORD")
            
            # Clean up URI format (remove spaces around colon)
            uri = uri.replace(" :", ":").replace(": ", ":")
            
            # Debug logging
            logger.info(f"Attempting to connect to Neo4j with URI: {uri}, Username: {username}")
            
            if not password:
                logger.error("NEO4J_PASSWORD environment variable is required")
                return None
            
            driver = GraphDatabase.driver(uri, auth=(username, password))
            # Test connection
            with driver.session() as session:
                session.run("RETURN 1")
            logger.info(f"Successfully connected to Neo4j at {uri}")
        except (ServiceUnavailable, AuthError) as e:
            logger.error(f"Failed to connect to Neo4j: {e}")
            logger.warning("Database operations will fail until connection is established")
        except Exception as e:
            logger.error(f"Unexpected error connecting to Neo4j: {e}")
    return driver

@mcp.tool()
def create_node(label: str, properties: Optional[Dict[str, Any]] = None) -> str:
    """Create a new node in the Neo4j database.
    
    Args:
        label: The label for the node
        properties: Optional dictionary of properties for the node
    
    Returns:
        Success message with node details
    """
    try:
        driver = get_neo4j_driver()
        if not driver:
            return "Error: No database connection available"
            
        with driver.session() as session:
            if properties:
                # Create node with properties
                props_str = ", ".join([f"{k}: ${k}" for k in properties.keys()])
                query = f"CREATE (n:{label} {{{props_str}}}) RETURN n"
                result = session.run(query, **properties)
            else:
                # Create node without properties
                query = f"CREATE (n:{label}) RETURN n"
                result = session.run(query)
            
            node = result.single()["n"]
            return f"Created node with label '{label}' and ID {node.id}: {dict(node)}"
    except Exception as e:
        logger.error(f"Error creating node: {e}")
        return f"Error creating node: {str(e)}"

@mcp.tool()
def create_edge(from_node_id: int, to_node_id: int, relationship_type: str, properties: Optional[Dict[str, Any]] = None) -> str:
    """Create a relationship between two nodes.
    
    Args:
        from_node_id: ID of the source node
        to_node_id: ID of the target node
        relationship_type: Type of the relationship
        properties: Optional dictionary of properties for the relationship
    
    Returns:
        Success message with relationship details
    """
    try:
        driver = get_neo4j_driver()
        if not driver:
            return "Error: No database connection available"
            
        with driver.session() as session:
            if properties:
                props_str = ", ".join([f"{k}: ${k}" for k in properties.keys()])
                query = f"""
                MATCH (a), (b)
                WHERE id(a) = $from_id AND id(b) = $to_id
                CREATE (a)-[r:{relationship_type} {{{props_str}}}]->(b)
                RETURN r
                """
                params = {"from_id": from_node_id, "to_id": to_node_id, **properties}
            else:
                query = f"""
                MATCH (a), (b)
                WHERE id(a) = $from_id AND id(b) = $to_id
                CREATE (a)-[r:{relationship_type}]->(b)
                RETURN r
                """
                params = {"from_id": from_node_id, "to_id": to_node_id}
            
            result = session.run(query, params)
            relationship = result.single()["r"]
            return f"Created relationship '{relationship_type}' between nodes {from_node_id} and {to_node_id}: {dict(relationship)}"
    except Exception as e:
        logger.error(f"Error creating relationship: {e}")
        return f"Error creating relationship: {str(e)}"

@mcp.tool()
def query_nodes(label: Optional[str] = None, properties: Optional[Dict[str, Any]] = None, limit: int = 10) -> str:
    """Query nodes from the Neo4j database.
    
    Args:
        label: Optional label to filter nodes
        properties: Optional properties to match
        limit: Maximum number of nodes to return (default: 10)
    
    Returns:
        JSON string containing the query results
    """
    try:
        driver = get_neo4j_driver()
        if not driver:
            return "Error: No database connection available"
            
        with driver.session() as session:
            if label and properties:
                where_clauses = [f"n.{k} = ${k}" for k in properties.keys()]
                where_str = " AND ".join(where_clauses)
                query = f"MATCH (n:{label}) WHERE {where_str} RETURN n LIMIT $limit"
                params = {**properties, "limit": limit}
            elif label:
                query = f"MATCH (n:{label}) RETURN n LIMIT $limit"
                params = {"limit": limit}
            elif properties:
                where_clauses = [f"n.{k} = ${k}" for k in properties.keys()]
                where_str = " AND ".join(where_clauses)
                query = f"MATCH (n) WHERE {where_str} RETURN n LIMIT $limit"
                params = {**properties, "limit": limit}
            else:
                query = "MATCH (n) RETURN n LIMIT $limit"
                params = {"limit": limit}
            
            result = session.run(query, params)
            nodes = []
            for record in result:
                node = record["n"]
                nodes.append({
                    "id": node.id,
                    "labels": list(node.labels),
                    "properties": dict(node)
                })
            
            return json.dumps({"nodes": nodes, "count": len(nodes)}, indent=2)
    except Exception as e:
        logger.error(f"Error querying nodes: {e}")
        return f"Error querying nodes: {str(e)}"

@mcp.tool()
def query_edges(relationship_type: Optional[str] = None, limit: int = 10) -> str:
    """Query relationships from the Neo4j database.
    
    Args:
        relationship_type: Optional relationship type to filter
        limit: Maximum number of relationships to return (default: 10)
    
    Returns:
        JSON string containing the query results
    """
    try:
        driver = get_neo4j_driver()
        if not driver:
            return "Error: No database connection available"
            
        with driver.session() as session:
            if relationship_type:
                query = f"MATCH (a)-[r:{relationship_type}]->(b) RETURN a, r, b LIMIT $limit"
            else:
                query = "MATCH (a)-[r]->(b) RETURN a, r, b LIMIT $limit"
            
            result = session.run(query, {"limit": limit})
            relationships = []
            for record in result:
                start_node = record["a"]
                relationship = record["r"]
                end_node = record["b"]
                
                relationships.append({
                    "id": relationship.id,
                    "type": relationship.type,
                    "properties": dict(relationship),
                    "start_node": {
                        "id": start_node.id,
                        "labels": list(start_node.labels),
                        "properties": dict(start_node)
                    },
                    "end_node": {
                        "id": end_node.id,
                        "labels": list(end_node.labels),
                        "properties": dict(end_node)
                    }
                })
            
            return json.dumps({"relationships": relationships, "count": len(relationships)}, indent=2)
    except Exception as e:
        logger.error(f"Error querying relationships: {e}")
        return f"Error querying relationships: {str(e)}"

@mcp.tool()
def delete_node(node_id: int) -> str:
    """Delete a node from the Neo4j database.
    
    Args:
        node_id: ID of the node to delete
    
    Returns:
        Success message
    """
    try:
        driver = get_neo4j_driver()
        if not driver:
            return "Error: No database connection available"
            
        with driver.session() as session:
            # First check if node exists
            check_query = "MATCH (n) WHERE id(n) = $node_id RETURN n"
            check_result = session.run(check_query, {"node_id": node_id})
            if not check_result.single():
                return f"Node with ID {node_id} not found"
            
            # Delete the node and all its relationships
            delete_query = "MATCH (n) WHERE id(n) = $node_id DETACH DELETE n"
            session.run(delete_query, {"node_id": node_id})
            return f"Successfully deleted node with ID {node_id}"
    except Exception as e:
        logger.error(f"Error deleting node: {e}")
        return f"Error deleting node: {str(e)}"

@mcp.tool()
def delete_edge(relationship_id: int) -> str:
    """Delete a relationship from the Neo4j database.
    
    Args:
        relationship_id: ID of the relationship to delete
    
    Returns:
        Success message
    """
    try:
        driver = get_neo4j_driver()
        if not driver:
            return "Error: No database connection available"
            
        with driver.session() as session:
            # First check if relationship exists
            check_query = "MATCH ()-[r]->() WHERE id(r) = $rel_id RETURN r"
            check_result = session.run(check_query, {"rel_id": relationship_id})
            if not check_result.single():
                return f"Relationship with ID {relationship_id} not found"
            
            # Delete the relationship
            delete_query = "MATCH ()-[r]->() WHERE id(r) = $rel_id DELETE r"
            session.run(delete_query, {"rel_id": relationship_id})
            return f"Successfully deleted relationship with ID {relationship_id}"
    except Exception as e:
        logger.error(f"Error deleting relationship: {e}")
        return f"Error deleting relationship: {str(e)}"

@mcp.tool()
def execute_cypher(query: str, parameters: Optional[Dict[str, Any]] = None) -> str:
    """Execute a custom Cypher query.
    
    Args:
        query: The Cypher query to execute
        parameters: Optional parameters for the query
    
    Returns:
        JSON string containing the query results
    """
    try:
        driver = get_neo4j_driver()
        if not driver:
            return "Error: No database connection available"
            
        with driver.session() as session:
            result = session.run(query, parameters or {})
            records = []
            for record in result:
                record_dict = {}
                for key in record.keys():
                    value = record[key]
                    if hasattr(value, 'id'):  # Node or Relationship
                        if hasattr(value, 'labels'):  # Node
                            record_dict[key] = {
                                "id": value.id,
                                "labels": list(value.labels),
                                "properties": dict(value)
                            }
                        else:  # Relationship
                            record_dict[key] = {
                                "id": value.id,
                                "type": value.type,
                                "properties": dict(value)
                            }
                    else:
                        record_dict[key] = value
                records.append(record_dict)
            
            return json.dumps({"records": records, "count": len(records)}, indent=2)
    except Exception as e:
        logger.error(f"Error executing Cypher query: {e}")
        return f"Error executing Cypher query: {str(e)}"

def main():
    """Main function to run the MCP server"""
    # Debug: Print environment variables
    logger.info("=== Neo4j MCP Server Starting ===")
    logger.info(f"NEO4J_URI: {os.getenv('NEO4J_URI', 'Not set')}")
    logger.info(f"NEO4J_USERNAME: {os.getenv('NEO4J_USERNAME', 'Not set')}")
    logger.info(f"NEO4J_PASSWORD: {'Set' if os.getenv('NEO4J_PASSWORD') else 'Not set'}")
    logger.info(f"LOG_LEVEL: {os.getenv('LOG_LEVEL', 'Not set')}")
    
    # Initialize Neo4j connection
    get_neo4j_driver()
    
    # Run the server
    mcp.run()

if __name__ == "__main__":
    main()
